import os
import math
import torch
import torchvision
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np



class feature_map_visualize:
    subplot_limit_x = 8
    subplot_limit_y = 8
    feature_map_resize= 250
    upsample = torch.nn.Upsample(size=(feature_map_resize, feature_map_resize), mode='bilinear', align_corners=True)
    # upsample = torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
    def __init__(self, model: torch.nn.Module,  names=[], save_dir='./') -> None:
        self.save_dir = save_dir
        self.model = model
        self.name_dicts = {}
        for i, name in enumerate(names):
            self.name_dicts[name] = 'feat{}'.format(i)
        self.get_layers = torchvision.models._utils.IntermediateLayerGetter(self.model, self.name_dicts)            
        self.input = None
        self.names_children = self.model.named_children()
        print('model\'s name children:')
        print([name[0] for name in self.names_children])
        self.all_layers = model.children()

    def fit_input(self, input: torch.tensor):
        print('Getting output')
        self.input = input

    def fit_name_list(self, names):
        self.name_dicts = {}
        for i, name in enumerate(names):
            self.name_dicts[name] = 'feat{}'.format(i)
        self.get_layers = torchvision.models._utils.IntermediateLayerGetter(self.model, self.name_dicts)            

    def visualize(self, index):
        print('visualize ...')
        if self.input is None:
            print('Input is none')
            return
        
        if isinstance(index, str):
            self.visualize_with_name(index)

    def visualize_with_name(self, name):
        self.fit_name_list([name])
        out = self.get_layers(self.input)
        print('get feature map from {}'.format(name))
        subplot_total = self.subplot_limit_x * self.subplot_limit_y
        for k, v in out.items():
            v_list = []
            if isinstance(v, tuple) or isinstance(v, list):
                v_list = v
            else:
                v_list = [v]
            for i, feature_map in enumerate(v_list):
                # print(feature_map.shape)
                feature_map = self.upsample(feature_map).squeeze(0).detach().cpu()
                feature_chunks = torch.chunk(feature_map, math.ceil(feature_map.shape[0] / subplot_total), 0)
                for j, feature_chunk in enumerate(feature_chunks):
                    # print(feature_chunk.shape)
                    figure_name = '{}_{}_{}'.format(k, i, j)
                    # feature_chunk = self.upsample(feature_chunk)
                    plt.figure(figure_name)
                    for n, feature in enumerate(feature_chunk):
                        # print(feature.shape)
                        plt.subplot(self.subplot_limit_x, self.subplot_limit_y, n + 1)
                        plt.imshow(feature, cmap='gray')
                        plt.axis('off')
                    plt.savefig(os.path.join(self.save_dir, '{}.png'.format(figure_name)))